%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%                         Transitional Dynamics                           %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function [InfoMatrix, u_func] = Trans_FP(p, InfoMatrix, u_func, u_Deltagrid, F1, FNC, Deltagrid)

% ---------------------------- Description -------------------------------- 

    % This script uses the equilibrium conditions to update a given
    % production share path. It repeats the process until the difference
    % between two iterations is small.

%--------------------------------------------------------------------------
%                             Preliminaries                               %
%--------------------------------------------------------------------------

    % Recover the objects that are encapsulated in the Info Matrix.
    eta_vector = InfoMatrix(1,:);
    LFN_t = InfoMatrix(2,:);
    gN_vector = InfoMatrix(3,:);
    prod_share_vector = InfoMatrix(4,:);
    M_vector = InfoMatrix(5,:);
    Lam_vector = InfoMatrix(6,:);
    gLam = InfoMatrix(7,:);

    % Recover counterfactual misallocation as the final terms of the vector
    Lam_counter = Lam_vector(end);
    M_counter = M_vector(end);

    % Determine the number of periods after the shock takes place
    TimePeriods = p.Tperiods;

    % Step size implied in the gap grid for the b function
    u_Delta_step = u_Deltagrid(2) - u_Deltagrid(1);

    % Number of points in Delta grid
    Deltagrid_pts = length(Deltagrid);

    % ---------------------------------------------------------------------

    % Initial guess for the production share path
    prod_share_new = prod_share_vector;

    % Variable that will be the stopping condition for the algorithm
    Stop = 0;

    % Determines the tolerance for the convergence of the algorithm
    tol = 2.5 * 1e-4;

    % Starts counting the number of iterations
    Iter = 1;

    % Determines whether we start with an almost solution guess
    AlmostSol = 0;

    % Determines the maximum number of iterations
    MaxIter = AlmostSol*20 + (1-AlmostSol)*2020;

    % Variable that determines whether misallocation terms are updated
    Misalloc_Update = 0;

    % Stop conditions are established in the end of while loop
    while Stop == 0 && Iter < MaxIter

%--------------------------------------------------------------------------
%                  1. Computing the variety growth rate                   %
%--------------------------------------------------------------------------

    % Picks previous production share path
    prod_share_old = prod_share_new;

    % For each time period, we compute the variety growth rate
    for t = 2:TimePeriods

        % Sets the market clearing condition as function of gN - see Equation (RP-45)
        root_func = @(x) (LFN_t(t-1)*exp((eta_vector(t)-x)*p.yearlength)*(1-prod_share_old(t))-((x+p.d0)/(1-p.alpha)-(p.zeta-1)/p.zeta*p.xstar)/p.phie);
            
        % Use fzero to search for the root of the proposed function - see Equation (RP-45)
        gN_vector(t) = fzero(root_func, gN_vector(t-1));

        % Update current Labor Force to Variety ratio - see Equation (RP-45)
        LFN_t(t) = LFN_t(t-1) * exp((eta_vector(t)-gN_vector(t))*p.yearlength);

    end

%--------------------------------------------------------------------------
%                2. Computing the Misallocation Terms                     %
%--------------------------------------------------------------------------

    % Conditions when misallocation terms will be updated 
    if Misalloc_Update == 1 || Iter == 1

        % Creative destruction rate path  - see Equation (RP-46)
        tau_vector = p.alpha/(1-p.alpha)*(gN_vector + p.d0);

        % Average efficiency growth rate path - see Equation (RP-46)
        gQ_vector = p.In + (p.qbar - 1)*(gN_vector + p.d0)/((1-p.alpha)*(p.s-1));

        % Initial conditions are the distributions on the calibrated BGP
        Current_Fc = F1;
        Current_FNC = FNC;

        % Determine a final period for updating, by which SS should arrive
        FinalPeriod = (2100-1980)/p.yearlength;

        % Find the index associated with the year of 2020
        TimeCut = (2020-1980)/p.yearlength;

        % Updates will only occur before the FinalYear period...
        for t = 2:FinalPeriod

        % ------------------------ No competitors -------------------------

        % Compute the step size implied in the productivity grid
        delta_q = p.qgrid(2) - p.qgrid(1);

        % Obtain the Gamma distribution - see Equation (RP-47)
        Gammadist = max(0, 1 - (p.wmin*exp(-p.qgrid)).^p.varrho);

        % Derivative wrt q over distribution in previous time period  - see Equation (RP-47)
        dervec = diff(Current_FNC);
        dfdq = [dervec(1), dervec] / delta_q;

        % Solve the differential equation  - see Equation (RP-47)
        dfdt_NC = -(p.In-gQ_vector(t-1))*dfdq - (tau_vector(t-1)+p.d0+gN_vector(t-1))*Current_FNC + (1-p.alpha)/p.alpha*tau_vector(t-1)*Gammadist;
        
        % Obtain the distribution over the next time period  - see Equation (RP-47)
        new_FNC = Current_FNC + dfdt_NC*p.yearlength;

        % ----------------------------- Competitors ------------------------------

        % Preallocate a matrix to store the distribution over the next period
        new_Fc = zeros(Deltagrid_pts, p.n3);

        % Compute the (lagged) terms that do not depend on Delta
        fixedterm = interp1(p.qgrid, Current_FNC, p.qgrid - log(p.l), "linear", 0) + interp1(p.qgrid, Current_Fc(end,:), p.qgrid - log(p.l), "linear", 0);
        
        % For each possible productivity gap value...
        for d = 2:Deltagrid_pts

            % Derivative with respect to productivity  - see Equation (RP-48)
            dfdq = [0, diff(Current_Fc(d, :))/delta_q];

            % Derivative with respect to Delta  - see Equation (RP-48)
            dfDelta = (Current_Fc(d,:)-Current_Fc(d-1,:))/(Deltagrid(d)-Deltagrid(d-1));

            % Derivative with respect to time  - see Equation (RP-48)
            dfdt = -(p.In-gQ_vector(t-1))*dfdq - p.In*Deltagrid(d)*dfDelta - (tau_vector(t-1)+p.d0+gN_vector(t-1))*Current_Fc(d,:) + tau_vector(t-1)*fixedterm;

            % Compute new distribution - see Equation (RP-48)
            new_Fc(d, :) = Current_Fc(d, :) + dfdt*p.yearlength;
        
        end

        % --------------- Updating Distributions --------------------------

        % Obtain the distribution for the current time period
        Current_FNC = new_FNC;

        % Obtain the distribution for the current time period
        Current_Fc = new_Fc;

        % -------------------------- Densities ----------------------------

        % Determines the grid for q that is implied in the grid for log(q)
        Q_endog_grid = exp(p.qgrid);

        % Obtains density for varieties without competitors defined over log(q)
        FNCdens = zeros(1, p.n3);
        FNCdens(1:end-1) = diff(Current_FNC)/delta_q;

        % Transforms this density to be defined over q
        FNCdens_endog = FNCdens ./ Q_endog_grid;
        FNCdens_endog(isnan(FNCdens_endog)) = 0;

        % Obtains density for varieties with competitors defined over log(q)
        FdensC = zeros(Deltagrid_pts, p.n3);
        FdensC(:, 1:end-1) = diff(Current_Fc, 1, 2)/delta_q;

        % Transforms this density to be defined over q
        FCdens_endog = FdensC ./ repmat(Q_endog_grid, Deltagrid_pts, 1);

        % Also, for every value of productivity...
        for q = 1:p.n3

            % ...we need to take differences over the gap dimension
            FCdens_endog(:, q) = gradient(FCdens_endog(:, q), Deltagrid);

        end

        % ------------------------ Numerical Issues -----------------------

        % Compute exponential terms that enter the integrals - see Equation (RP-21)
        exp_term = repmat(Q_endog_grid.^(p.s-1), Deltagrid_pts, 1);

        % Compute the adjustment factor for the integrals  - see Equation (RP-21)
        Adj_NC = trapz(Q_endog_grid, FNCdens_endog .* Q_endog_grid.^(p.s-1));
        Adj_C = trapz(Q_endog_grid, trapz(Deltagrid, FCdens_endog.*exp_term, 1));
        Adj_fac = Adj_NC + Adj_C;

        % Constant that shifts the grid by making Equation (RP-21) hold
        Adj_constant = -log(Adj_fac)/(p.s-1);

        % Shifts the productivity grid accordingly
        qgrid_adj = p.qgrid + Adj_constant;

        % -------------------- Repeat previous steps ----------------------

        % Determines the grid for q that is implied in the grid for log(q)
        Q_endog_grid = exp(qgrid_adj);

        % Obtains density for varieties without competitors defined over log(q)
        FNCdens = zeros(1, p.n3);
        FNCdens(1:end-1) = diff(Current_FNC)/delta_q;

        % Transforms this density to be defined over q
        FNCdens_endog = FNCdens ./ Q_endog_grid;
        FNCdens_endog(isnan(FNCdens_endog)) = 0;

        % Obtains density for varieties with competitors defined over log(q)
        FdensC = zeros(Deltagrid_pts, p.n3);
        FdensC(:, 1:end-1) = diff(Current_Fc, 1, 2)/delta_q;
    
        % Transforms this density to be defined over q
        FCdens_endog = FdensC ./ repmat(Q_endog_grid, Deltagrid_pts, 1);

        % Also, for every value of productivity...
        for q = 1:p.n3

            % ...we need to take differences over the gap dimension
            FCdens_endog(:, q) = gradient(FCdens_endog(:, q), Deltagrid);

        end

        % ---------------------- Misallocation Terms ----------------------

        % Compute the mark-up and exponential terms that enter the integrals - see Equations (RP-13) and (RP-14)
        MarkupTerm = repmat(min(p.mu, Deltagrid),p.n3,1)';
        exp_term = repmat(Q_endog_grid.^(p.s-1), Deltagrid_pts, 1);

        % Compute the integral term associated with no competitors - see Equations (RP-13) and (RP-14)
        Int_NC = trapz(Q_endog_grid, FNCdens_endog .* Q_endog_grid.^(p.s-1));
        
        % Compute the first integrals associated varieties with competitors - see Equations (RP-13) and (RP-14)
        integrand = FCdens_endog.*exp_term.*MarkupTerm.^(1-p.s);
        Int_C1 = trapz(Q_endog_grid, trapz(Deltagrid, integrand, 1));

        % Compute the second integrals associated varieties with competitors - see Equations (RP-13) and (RP-14)
        integrand = FCdens_endog.*exp_term.*MarkupTerm.^(-p.s);
        Int_C2 = trapz(Q_endog_grid, trapz(Deltagrid, integrand, 1));

        % Compute the misallocation terms - see Equations (RP-13) and (RP-14)
        Lam_vector(t) = (Int_C2 + p.mu^(-p.s)*Int_NC)/(Int_C1 + p.mu^(1-p.s)*Int_NC);
        M_vector(t) = (Int_C1 + p.mu^(1-p.s)*Int_NC)^(1/(p.s-1))/Lam_vector(t);
        
        end

        % --------------------- Smooth Lambda path ------------------------

        % Smooth the Lambda path between TimeCut = 2020 and FinalPeriod = 2100
        [aux, ~] = hpfilter(Lam_vector(TimeCut+1:FinalPeriod), Smoothing = 1e8); 
        Lam_vector(TimeCut+1:FinalPeriod) = aux;

        % Make sure the two pieces of the path connect appropriately in 2020
        Lam_vector(TimeCut+1:FinalPeriod) = Lam_vector(TimeCut+1:FinalPeriod) - (Lam_vector(TimeCut+1)-Lam_vector(TimeCut));

        % Compute the adjustment needed at FinalPeriod = 2100 to reach new BGP Lambda
        Delta = (Lam_counter - Lam_vector(FinalPeriod));

        % Shift the path for Lambda accordingly
        Lam_vector(TimeCut+1:FinalPeriod) = Lam_vector(TimeCut+1:FinalPeriod)+(p.TimeGrid(TimeCut+1:FinalPeriod)-p.TimeGrid(TimeCut+1))/(p.TimeGrid(FinalPeriod)-p.TimeGrid(TimeCut+1)).*Delta;

        % After FinalPeriod = 2100, the path for Lambda is constant
        Lam_vector(FinalPeriod+1:end) = Lam_vector(FinalPeriod);

        % Compute the growth rate of Lambda at the time period unit - see Equation (RP-49)
        gLam = [0, log(Lam_vector(2:end)./Lam_vector(1:end-1))];

        % Sums over the last 12 months (when possible) for annual rates - see Equation (RP-49)
        gLam = movmean(gLam, [(1/p.yearlength-1) 0]) .* min(linspace(1, TimePeriods, TimePeriods), 1/p.yearlength);

        % Ensures smoothness of gLam at FinalPeriod = 2100 - see Equation (RP-49)
        gLam(FinalPeriod+1:end) = gLam(FinalPeriod);

        % ----------------------- Smooth M path ---------------------------

        % Smooth the M path between TimeCut = 2020 and FinalPeriod = 2100
        [aux, ~] = hpfilter(M_vector(TimeCut+1:FinalPeriod), Smoothing = 1e8); 
        M_vector(TimeCut+1:FinalPeriod) = aux;

        % Make sure the two pieces of the path connect appropriately in 2020
        M_vector(TimeCut+1:FinalPeriod) = M_vector(TimeCut+1:FinalPeriod) - (M_vector(TimeCut+1)-M_vector(TimeCut));

        % Compute the adjustment needed at FinalPeriod = 2100 to reach new BGP M
        Delta = (M_counter - M_vector(FinalPeriod));

        % Shift the path for M accordingly
        M_vector(TimeCut+1:FinalPeriod) = M_vector(TimeCut+1:FinalPeriod)+(p.TimeGrid(TimeCut+1:FinalPeriod)-p.TimeGrid(TimeCut+1))/(p.TimeGrid(FinalPeriod)-p.TimeGrid(TimeCut+1)).*Delta;

        % After FinalPeriod = 2100, the path for M is constant
        M_vector(FinalPeriod+1:end) = M_vector(FinalPeriod);

    end

%--------------------------------------------------------------------------
%                  3. Computing the b_t(Delta) function                   %
%--------------------------------------------------------------------------
    
    % Create the function h(Delta) according to its specification - see Section (RP-3.1)
    h_func = @(d) (min(p.s/(p.s-1),d)-1)./min(p.s/(p.s-1),d).^p.s;
   
    % For each time period...
    for t = 2:TimePeriods

        % ... compute the derivative wrt Delta over the previous column - see Equation (RP-44)
        dudDelta = [diff(u_func(:,t-1)) ./ u_Delta_step; 0];

        % ... solve the differential equation for the time derivative  - see Equation (RP-44)
        dudt = h_func(u_Deltagrid')/(M_vector(t-1)^(p.s-1)*Lam_vector(t-1)^p.s) + p.In * u_Deltagrid' .* dudDelta * p.yearlength;
        dudt = dudt - (p.rho - gLam(t-1) - eta_vector(t-1)+(p.d0+gN_vector(t-1))/(1-p.alpha)*p.qbar) .* u_func(:, t-1) * p.yearlength;

        % ... obtain the function b_t(Delta) for the next period - see Equation (RP-44)
        u_func(:, t) = u_func(:, t-1) + dudt;

    end

%--------------------------------------------------------------------------
%               4. Update production share via free entry                 %
%--------------------------------------------------------------------------

    % Evaluate the function b_t(Delta) and derivative at Delta = lambda.
    ulambda = u_func(1, :);
    ulambda_d = [0, diff(ulambda)*p.yearlength];

    % Evaluate the function b_t(Delta) and derivative at Delta = CES.
    uCES =  u_func(end, :);
    uCES_d = [0, diff(uCES)*p.yearlength];

    % Computes the (forward) growth rates of share - see Footnote (RP-12)
    gShare = [0, log(prod_share_old(3:end)./prod_share_old(2:end-1)), 0];
    gShare = movmean(gShare, [0 (1/p.yearlength-1)]) .* fliplr(min(linspace(1, TimePeriods, TimePeriods), 1/p.yearlength));

    % Restore zero growth rate on the initial period - consistent w/ BGP
    gShare(1) = 0;

    % Compute the terms associated with b(lambda) and b(mu) - see Equation (RP-43)
    lambda_term = p.alpha*p.l^(p.s-1)*(ulambda.*(p.rho-gLam-eta_vector+(gN_vector+p.d0)/(1-p.alpha))*p.yearlength - ulambda_d).*p.phie.*LFN_t;
    CES_term = (1-p.alpha)*p.wbar*(uCES.*(p.rho-gLam-eta_vector+(gN_vector+p.d0)/(1-p.alpha))*p.yearlength - uCES_d).*p.phie.*LFN_t;

    % Use free entry to obtain the updated production share - see Equation (RP-43)
    prod_share_new = (p.rho + gShare - gLam + (p.alpha*gN_vector+p.d0)/(1-p.alpha)  - (p.zeta-1)/p.zeta*p.xstar)./(lambda_term+CES_term);
   
%--------------------------------------------------------------------------
%                         5. Close the Loop                               %
%--------------------------------------------------------------------------

    % Measure the difference between the production shares
    diff_shares = max(abs(prod_share_new - prod_share_old), [], 'all');

    % If error is small and misallocation has been updated...
    if diff_shares < tol && Misalloc_Update == 1

        % ... the code ends successfully
        Stop = 1;

    % If error is small on the first iteration... 
    elseif diff_shares < tol && Iter == 1

        % ... the code ends successfully because misallocation is always updated on first iteration
        Stop = 1;
    end

    % ---------------------------------------------------------------------

    % There will only be misallocation updates at every iteration after iteration 2000...
    if rem(Iter + 1, 1000) == 0 || (Iter + 1) >= 2000
       Misalloc_Update = 1;
    else
       Misalloc_Update = 0;
    end

    % ... but if AlmostSol = 1, then updates occur in every iteration
    if AlmostSol == 1
        Misalloc_Update = 1;
    end

    % Determines frequency with which track of the algorithm is printed
    Print_Freq = 500*(1-Misalloc_Update) + Misalloc_Update;

    % Determines the specific iterations at which information is printed
    if Iter == 1 || rem(Iter, Print_Freq) == 0

        % ...it prints the current difference between paths
        fprintf('     Iteration %d: difference in production share paths is %.7f. \n', Iter, diff_shares);


        % ... and shows a plot of the current situation
        close
        figure('Name','Transitional Dynamics - Update');
        plot(prod_share_old, 'DisplayName','Old Share', 'LineWidth', 1)
        hold on
        plot(prod_share_new, 'DisplayName','New Share', 'LineWidth', 1)
        legend('Location', 'southeast')
        drawnow;

    end

    % We use a very conservative weight to update production share
    weight = 0.995;

    % New path is a convex combination of the two paths
    prod_share_new = weight*prod_share_old + (1-weight)*prod_share_new;

    % Update the iteration counter
    Iter = Iter + 1;

    end

    % Assign the resulting production share path to usual vector
    prod_share_vector = prod_share_old;

    % Stores all relevant quantities to the InfoMatrix
    InfoMatrix = [eta_vector; LFN_t; gN_vector; prod_share_vector; M_vector; Lam_vector; gLam];
    close

end
    


